package io.gitee.chearnee.fw.log.sql.core;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.slf4j.MDC;

import java.sql.Connection;


/**
 * sql中添加当前traceId的拦截器
 *
 * @author lk
 */
@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
@Slf4j
public class TraceIdInterceptor implements Interceptor {

    /**
     * traceId
     */
    private final String traceId;

    public TraceIdInterceptor(String traceId) {

        this.traceId = traceId;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        String s = MDC.get(traceId);
        if (s == null || s.trim().isEmpty()) {
            return invocation.proceed();
        }
        Object target = invocation.getTarget();
        if (target instanceof StatementHandler statementHandler) {
            BoundSql boundSql = statementHandler.getBoundSql();
            MetaObject metaObject = SystemMetaObject.forObject(boundSql);
            String sql = boundSql.getSql();
            String newSql = sql + "/*【traceId:" + s + "】*/";
            metaObject.setValue("sql", newSql);
        }
        return invocation.proceed();
    }


    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor || target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }
}